import threading
import time
from queue import Empty, Full, Queue
from typing import Callable

import numpy as np

from cion.emccd.camera import AcquisitionMode, DRV_NO_NEW_DATA, EmccdCamera, ReadMode, ReturnCodeError, ShutterMode, \
    ShutterTTLSignal, create_uint32_buffer


def capture_video(
        camera: EmccdCamera,
        callback: Callable[[np.ndarray], None],
        **kwargs,
) -> 'VideoCapturing':
    return VideoCapturing(camera, callback, **kwargs)


class VideoCapturing:
    def __init__(self,
            camera: EmccdCamera,
            callback: Callable[[np.ndarray], None],
            **kwargs):
        self._queue = Queue(1)
        thread = threading.Thread(
            daemon=True,
            target=_video_thread_func,
            args=(camera, self._queue, callback),
            kwargs=kwargs)
        thread.start()
        self._thread = thread

    def __del__(self):
        self.stop()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()

    def stop(self):
        queue = self._queue
        if queue is None:
            return

        try:
            queue.put_nowait(True)
        except Full:
            pass

        thread = self._thread
        if thread is not None:
            thread.join()


def _video_thread_func(
        camera: EmccdCamera,
        queue: Queue,
        callback: Callable[[np.ndarray], None],
        **kwargs
):
    # temperature
    camera.set_temperature(-80)
    camera.set_cooler(True)

    # shutter
    camera.set_shutter(ShutterTTLSignal.LOW, ShutterMode.PERMANENTLY_OPEN, 27, 27)

    # baseline clamp
    if camera.support_baseline_clamp:
        camera.set_baseline_clamp(True)
    if camera.support_baseline_offset:
        camera.set_baseline_offset(0)

    # set shift speeds & ad channels
    camera.set_hs_speed(0, 0)
    camera.set_vs_speed(4)
    # camera.dll.SetVSAmplitude(0)
    # camera.dll.SetFrameTransferMode(0)

    # pre-amp gain
    camera.set_pre_amp_gain(2)

    # em & em gain
    camera.set_em_advanced(False)
    camera.set_em_gain(200)

    # set read mode & set image
    camera.set_read_mode(ReadMode.IMAGE)
    img_w, img_h = camera.get_detector_size()
    camera.set_image(hbin=1, vbin=1, hstart=0, hend=img_w, vstart=0, vend=img_h)
    img_size = img_w * img_h

    # set acquisition mode
    camera.set_acquisition_mode(AcquisitionMode.RUN_TILL_ABORT)
    camera.set_exposure_time(0.1)
    camera.set_kinetic_cycle_time(0.1)
    exposure_time, accumulate_cycle_time, kinetic_cycle_time = camera.get_acquisition_times()
    # camera.dll.SetTriggerMode(0)

    # prepare buffer
    img_buffer = create_uint32_buffer(img_size)
    img_arr = np.frombuffer(img_buffer, dtype=np.uint32).reshape((img_h, img_w))

    # capture
    camera.start_acquisition()
    while True:
        try:
            queue.get_nowait()
            break
        except Empty:
            pass

        start_time = time.time()

        try:
            camera.get_latest_image(img_size, img_buffer)
            callback(img_arr)
        except ReturnCodeError as err:
            if err.rc != DRV_NO_NEW_DATA:
                raise err  # ignore DRV_NO_NEW_DATA

        try:
            print(f"temperature = {camera.get_temperature()}")
        except ReturnCodeError as err:
            print(f"temperature = ?? (rc={err.rc})")

        end_time = time.time()
        span_time = end_time - start_time
        if span_time < kinetic_cycle_time:
            time.sleep(kinetic_cycle_time - span_time)
    camera.abort_acquisition()
